import os
import scipy.io
import numpy as np
import matplotlib.pyplot as plt
from matplotlib.animation import FuncAnimation
from mpl_toolkits.mplot3d import Axes3D

def visualize_skeleton(mat_file_path):
    """可视化骨骼数据"""
    if not os.path.exists(mat_file_path):
        print(f"文件不存在: {mat_file_path}")
        return
    
    print(f"正在加载文件: {mat_file_path}")
    try:
        # 加载mat文件
        mat_data = scipy.io.loadmat(mat_file_path)
        
        # 假设数据存储在'data'变量中
        if 'data' not in mat_data:
            print("文件中没有找到'data'变量")
            return
            
        skeleton_data = mat_data['data']
        print(f"骨骼数据形状: {skeleton_data.shape}")
        
        # 获取帧数和关节点数
        num_frames, num_features = skeleton_data.shape
        num_joints = num_features // 2  # 假设每个关节有x,y两个坐标
        
        print(f"帧数: {num_frames}, 关节点数: {num_joints}")
        
        # 定义关节点之间的连接关系（根据通用骨骼模型近似）
        # 这里的连接关系需要根据实际关节点索引进行调整
        connections = []
        
        # 如果是25个关节点（MediaPipe标准）
        if num_joints == 25:
            # 定义面部连接
            face_connections = [(0, 1), (1, 2), (2, 3), (3, 4), (0, 5), (5, 6), (6, 7), (7, 8)]
            # 定义上半身连接
            upper_body = [(9, 10), (10, 11), (11, 12), (9, 13), (13, 14), (14, 15)]
            # 定义下半身连接
            lower_body = [(9, 16), (16, 17), (17, 18), (9, 19), (19, 20), (20, 21)]
            # 组合所有连接
            connections = face_connections + upper_body + lower_body
        
        # 创建交互式绘图
        plt.ion()
        
        # 创建2D图形
        fig = plt.figure(figsize=(12, 10))
        
        # 2D视图（左）
        ax1 = fig.add_subplot(121)
        ax1.set_title('2D骨骼视图')
        ax1.set_xlabel('X坐标')
        ax1.set_ylabel('Y坐标')
        
        # 3D视图（右）- 假设z坐标是0（如果没有z坐标）
        ax2 = fig.add_subplot(122, projection='3d')
        ax2.set_title('3D骨骼视图')
        ax2.set_xlabel('X坐标')
        ax2.set_ylabel('Y坐标')
        ax2.set_zlabel('帧数')
        
        # 创建滑块来选择帧
        from matplotlib.widgets import Slider
        slider_ax = plt.axes([0.2, 0.02, 0.6, 0.03])
        frame_slider = Slider(slider_ax, '帧数', 0, num_frames - 1, valinit=0, valstep=1)
        
        # 渲染选定帧的骨骼
        scatter1, = ax1.plot([], [], 'ro', ms=8)  # 关节点
        lines1 = []  # 连接线
        
        # 3D图中显示所有帧的轨迹
        scatter2 = ax2.scatter([], [], [], c='r', s=50)
        lines2 = []  # 3D连接线

        def update_frame(frame_idx):
            frame_idx = int(frame_idx)
            
            # 获取当前帧骨骼数据
            frame_data = skeleton_data[frame_idx]
            
            # 分离X和Y坐标
            x_coords = [frame_data[i] for i in range(0, len(frame_data), 2)]
            y_coords = [frame_data[i + 1] for i in range(0, len(frame_data), 2)]
            
            # 更新2D视图
            scatter1.set_data(x_coords, y_coords)
            ax1.set_xlim(min(x_coords) - 10, max(x_coords) + 10)
            ax1.set_ylim(min(y_coords) - 10, max(y_coords) + 10)
            
            # 清除之前的连接线
            for line in lines1:
                line.remove()
            lines1.clear()
            
            # 绘制新的连接线
            for start, end in connections:
                if start < len(x_coords) and end < len(x_coords):
                    line, = ax1.plot([x_coords[start], x_coords[end]], [y_coords[start], y_coords[end]], 'b-')
                    lines1.append(line)
            
            # 标记关节点编号
            for i, (x, y) in enumerate(zip(x_coords, y_coords)):
                ax1.text(x, y, str(i), fontsize=9)
            
            # 更新3D视图
            all_x = []
            all_y = []
            all_z = []
            colors = []
            
            # 显示所有帧的关节点，当前帧高亮显示
            for i in range(num_frames):
                frame_data_i = skeleton_data[i]
                x_i = [frame_data_i[j] for j in range(0, len(frame_data_i), 2)]
                y_i = [frame_data_i[j + 1] for j in range(0, len(frame_data_i), 2)]
                z_i = [i] * len(x_i)  # 使用帧索引作为z坐标
                
                all_x.extend(x_i)
                all_y.extend(y_i)
                all_z.extend(z_i)
                
                # 高亮显示当前帧
                if i == frame_idx:
                    colors.extend(['red'] * len(x_i))
                else:
                    colors.extend(['blue'] * len(x_i))
            
            ax2.cla()
            ax2.set_title('3D骨骼视图')
            ax2.set_xlabel('X坐标')
            ax2.set_ylabel('Y坐标')
            ax2.set_zlabel('帧数')
            ax2.scatter(all_x, all_y, all_z, c=colors, s=30, alpha=0.5)
            
            # 在3D视图中连接当前帧的关节点
            for start, end in connections:
                if start < len(x_coords) and end < len(x_coords):
                    ax2.plot([x_coords[start], x_coords[end]], 
                             [y_coords[start], y_coords[end]], 
                             [frame_idx, frame_idx], 'g-')
                    
            # 设置3D视图的轴范围
            ax2.set_xlim(min(all_x) - 10, max(all_x) + 10)
            ax2.set_ylim(min(all_y) - 10, max(all_y) + 10)
            ax2.set_zlim(0, num_frames)
                
            # 显示帧信息
            plt.suptitle(f"骨骼数据可视化 - 帧 {frame_idx+1}/{num_frames}")
            
            fig.canvas.draw_idle()
            
        # 将更新函数连接到滑块
        frame_slider.on_changed(update_frame)
        
        # 初始化第一帧
        update_frame(0)
        
        # 创建动画
        def animate():
            for i in range(num_frames):
                frame_slider.set_val(i)
                plt.pause(0.1)  # 控制动画速度
        
        # 创建动画按钮
        from matplotlib.widgets import Button
        button_ax = plt.axes([0.8, 0.02, 0.1, 0.04])
        button = Button(button_ax, '播放')
        button.on_clicked(lambda event: animate())
        
        plt.tight_layout()
        plt.subplots_adjust(bottom=0.2)  # 为滑块留出空间
        plt.show()
        
    except Exception as e:
        import traceback
        print(f"加载文件时出错: {e}")
        print(traceback.format_exc())

if __name__ == "__main__":
    import argparse
    
    parser = argparse.ArgumentParser(description="可视化骨骼数据工具")
    parser.add_argument("mat_file", help="要查看的mat文件的路径")
    
    args = parser.parse_args()
    visualize_skeleton(args.mat_file)
